import sys
import time
import random
import numpy as np
from graph_read import read_graph

class Tabu_color():
    def __init__(self,graph_filename,k):
        self.graph = read_graph(graph_filename)
        self.node_num = self.graph.shape[0]
        self.color_num = k
        self.extract_edges()
        self.tabu_list = {"move":[],"expiration":[]}
        self.tabu_respiration_best_score = 0
        self.cur_score = 0
        self.move = None
        self.tabutenure = 20
        self.check = False

    def init_solution(self):
        s_1 = np.zeros([self.node_num,self.color_num]).astype("int16")
        s_2 = np.zeros([self.node_num]).astype("int16")
                
        for i in range(self.node_num):
            s_2[i] = random.randint(0,self.color_num-1)
        for i in range(self.node_num):
            for j in range(self.node_num):
                if self.graph[i][j] == 1:
                    s_1[i][s_2[j]] += 1
        self.solution = (s_1,s_2)

    def extract_edges(self):
        self.edges = []
        for i in range(self.node_num):
            self.edges.append([])
            for j in range(self.node_num):
                if self.graph[i][j] > 0:
                    self.edges[i].append(j)

    def find_neighbors(self):
        self.neighbors = [] # we just store the move instead of specific list for cost
        neighbors_1 = [] # (n,i,j) and j is 0
        neighbors_2 = [] # (n,i,j) and j is less that i
        neighbors_3 = [] # (n,i,j) and j is others
        mask_ = self.solution[0][:,self.solution[1]] > 0
        for n in range(self.node_num):
            if self.solution[0][n][self.solution[1][n]] <= 0: continue
            for j in range(self.color_num):
                if self.solution[0][n][j] <= 0:
                    neighbors_1.append((n,self.solution[1][n],j))
                    continue
                if self.solution[0][n][j] <= self.solution[0][n][self.solution[1][n]]:
                    neighbors_2.append((n,self.solution[1][n],j))
                    continue
                neighbors_3.append((n,self.solution[1][n],j))
        self.neighbors.extend(neighbors_1)
        self.neighbors.extend(neighbors_2)
        self.neighbors.extend(neighbors_3)
        self.neighbors = self.neighbors[:100]

    def pick_solution(self):
        candidates = []
        for one in self.neighbors:
            val = self.eval_solution(one)
            candidates.append((one,val))
        candidates.sort(key=lambda one:one[1],reverse=True)
        candidates = candidates[:20]
        #candidates = [(one[0],one[1]+abs(one[1]+random.random())*random.random()) for one in candidates]
        #candidates.sort(key=lambda one:one[1],reverse=True)
        self.move = None
        if len(candidates) <= 0 or candidates[0][1] <= 0:
            self.check = True
        else:
            self.check = False
        for one in candidates:
            if (one[0][0],one[0][2]) not in self.tabu_list["move"]:
                self.move = one[0]
                self.cur_score += one[1]
                break

    def update_solution(self):
        if self.move is not None:
            edge_nodes = []
            for i in range(self.node_num):
                if self.graph[self.move[0],i] > 0:
                    self.solution[0][i][self.move[1]] -= 1
                    self.solution[0][i][self.move[2]] += 1
            self.solution[1][self.move[0]] = self.move[2]
            self.tabu_list["move"].append((self.move[0],self.move[1]))
            if len(self.tabu_list["move"]) > self.tabutenure:
                self.tabu_list["move"].pop(0)

    def eval_solution(self,move):
        val = 0
        val += self.solution[0][move[0]][move[1]]
        val -= self.solution[0][move[0]][move[2]]
        ret = np.bincount(self.solution[1][self.edges[move[0]]],minlength=self.node_num)
        val += ret[move[1]]
        val -= ret[move[2]]
        return val

    def check_solution(self,s):
        tmp = np.eye(self.color_num)[s[1]]
        tmp *= s[0]
        if np.sum(tmp) <= 0:
            return True
        return False

    def search(self,num):
        self.init_solution()
        i = 0
        start_time = time.time()
        while num < 0 or i < num:
            if self.check is True and self.check_solution(self.solution) is True:
                print("solution:%s" % str(t.solution))
                print("Find!")
                print("clock:%f" % time.clock())
                break
            if i% 1000 == 0:
                end_time = time.time()
                print("%d:not found,duration time:%f" % (i,end_time-start_time))
                if end_time - start_time > 3600:
                    break
            self.find_neighbors()
            self.pick_solution()
            self.update_solution()
            i+=1
        return self.solution
    
                    

if __name__ == "__main__":
    t = Tabu_color(sys.argv[1],int(sys.argv[2]))
    s_1 = time.time()
    t.search(int(sys.argv[3]))
    s_2 = time.time()
    print("s_2-s_1 time:%f" % (s_2-s_1))
